from yacs.config import CfgNode as CN

_C = CN()

_C.GPUS = [0]
_C.CUDA = True
_C.WORKERS = 16
_C.PRINT_FREQ = 5
_C.TEST_FREQ = 5
_C.LOSS = 'l2'
_C.OUTPUT_DIR = ''
_C.PRETRAIN = ''
_C.SYNCE_BN = False

# DATASET releated params
_C.DATASET = CN()
_C.DATASET.NAME = ''
_C.DATASET.ROOT = ''
_C.DATASET.ANNO_ROOT = ''
_C.DATASET.IMAGE_SIZE = [800, 800]

# MODEL releated params
_C.MODEL = CN()
_C.MODEL.NAME = 'model'
_C.MODEL.NUM_CLASSES = 0
_C.MODEL.BIAS_VALUE = -2.19
_C.MODEL.PRETRAIN = False
_C.MODEL.FPN_CHANNELS = []
_C.MODEL.RETURN_LAYERS = []


# TRAIN releated params
_C.TRAIN = CN()
_C.TRAIN.BATCH_SIZE_PER_GPU = 256
_C.TRAIN.MAX_EPOCH = 100
_C.TRAIN.OPTIMIZER = 'adam'
_C.TRAIN.LR = 1e-4
_C.TRAIN.LR_FACTOR = 0.1
_C.TRAIN.LR_STEP = [50, 100]
_C.TRAIN.WD = 5e-4
_C.TRAIN.MOENTUM = 0.9

_C.TRAIN.HEATMAP_WEIGHT = 1.0
_C.TRAIN.OFFSET_WEIGHT = 1.0
_C.TRAIN.WH_WEIGHT = 0.1
_C.TRAIN.LANDMARK_WEIGHT = 0.2

_C.TEST = CN()
_C.TEST.BATCH_SIZE_PER_GPU = 256



def update_config(cfg, opt):
    cfg.defrost()
    cfg.merge_from_file(opt.cfg)
    cfg.freeze()
